/*
 * model.h
 *
 *  Created on: Feb 2, 2012
 *      Author: kkb110
 */

#ifndef MODEL_H_
#define MODEL_H_

#include "common_includes.h"
#include "normal_distribution.h"

template<typename T, size_t N, size_t M>
class Model{
public:
	typedef Eigen::Matrix<T, M, 1> Shot;
	Eigen::Array<T, N, N> transition_probabilities;
	array<NormalDistribution<T, M> , N> shooting_distributions;

	Model(){

	}

	pair<vector<Shooter>, vector<Shot>> simulate(size_t n){
		vector<Shooter> shooters_order(n);
		vector<Shot> shots_order(n);

		random_device seed;
		mt19937 random_number_engine(seed());
		uniform_real_distribution<double> uniform_distribution_generator(0, 1); //Todo std::generate_canonical?

		T random_number = uniform_distribution_generator(random_number_engine);
		size_t shooter = N - 1;

		//random first shooter
		for (size_t i = 0; i < N; ++i){
			random_number -= 1.0 / N;
			if (random_number < 0){
				shooter = i;
				break;
			}
		}
		//fixed first shooter
		shooter = 0;

		shooters_order[0] = shooter;
		shots_order[0] = shooting_distributions[shooter].generate_points(1,
				random_number_engine)[0];

		for (size_t t = 1; t < n; ++t){
			random_number = uniform_distribution_generator(
					random_number_engine);
			shooter = N - 1;
			for (size_t i = 0; i < N; ++i){
				random_number -= transition_probabilities.row(
						shooters_order[t - 1])[i];
				if (random_number < 0){
					shooter = i;
					break;
				}
			}
			shooters_order[t] = shooter;
			shots_order[t] = shooting_distributions[shooter].generate_points(1,
					random_number_engine)[0];
		}
		return pair<vector<Shooter>, vector<Shot>>(shooters_order, shots_order);
	}

	inline vector<array<T, N>> forward_backward(
			const vector<Shot>& shots,
			const T& best_state_bias) const{
		vector<array<T, N>> result(shots.size());
		//Forward

		//initialization
		{
			T normalization = 0; // scaling factors
			for (unsigned i = 0; i < N; ++i){
				normalization += shooting_distributions[i].probability_density(
						shots[0]);
			}
			normalization = 1 / normalization;
			for (unsigned i = 0; i < N; ++i){
				result[0][i] = shooting_distributions[i].probability_density(
						shots[0]) * normalization;
			}
		}
//		initialization (fixed)
		{
			for (unsigned i = 0; i < N; ++i){
				result[0][i] = 0;
			}
			result[0][0] = 1;
		}

		//main loop
		for (unsigned t = 1; t < shots.size(); ++t){
			T normalization = 0; // scaling factors
			for (unsigned i = 0; i < N; ++i){
				result[t][i] = 0;
				for (unsigned j = 0; j < N; ++j){
					result[t][i] += result[t - 1][j]
							* transition_probabilities(j, i);
				}
				result[t][i] *= shooting_distributions[i].probability_density(
						shots[t]);
				normalization += result[t][i];
			}

			normalization = 1 / normalization;
			for (unsigned i = 0; i < N; ++i){
				result[t][i] *= normalization;
			}
		}

		//Backward
		vector<array<T, N>> backward(shots.size());

		//initialization
		for (unsigned i = 0; i < N; ++i){
			backward[shots.size() - 1][i] = 1.0 / N;
		}
		//main loop
		for (signed t = shots.size() - 2; 0 <= t; --t){
			array<T, N> temp;
			for (unsigned j = 0; j < N; ++j){
				temp[j] = backward[t + 1][j]
						* shooting_distributions[j].probability_density(
								shots[t + 1]);
			}
			T normalization = 0; // scaling factors
			for (unsigned i = 0; i < N; ++i){
				backward[t][i] = 0;
				for (unsigned j = 0; j < N; ++j){
					backward[t][i] += temp[j] * transition_probabilities(i, j);
				}
				normalization += backward[t][i];
			}
			normalization = 1 / normalization;
			for (unsigned i = 0; i < N; ++i){
				backward[t][i] *= normalization;
			}
		}

		//fb_combined
		for (unsigned t = 0; t < shots.size(); ++t){
			T normalization = 0;
			for (unsigned i = 0; i < N; ++i){
				result[t][i] *= backward[t][i];
				result[t][i] = pow(result[t][i], best_state_bias);
				normalization += result[t][i];
			}
			normalization = 1 / normalization;
			for (unsigned i = 0; i < N; ++i){
				result[t][i] *= normalization;
			}
		}

//		initialization (fixed)
		{
			for (unsigned i = 0; i < N; ++i){
				result[0][i] = 0;
			}
			result[0][0] = 1;
		}

		return result;
	}

	inline vector<array<T, N>> forward_backward_sequencewise(
			const vector<Shot>& shots,
			const T& best_path_bias) const{
		vector<array<T, N>> result(shots.size());
		//Forward

		//initialization
		{
			T normalization = 0; // scaling factors
			for (unsigned i = 0; i < N; ++i){
				normalization += pow(
						shooting_distributions[i].probability_density(shots[0]),
						best_path_bias);
			}
			normalization = 1 / normalization;
			for (unsigned i = 0; i < N; ++i){
				result[0][i] = pow(
						shooting_distributions[i].probability_density(shots[0]),
						best_path_bias) * normalization;
			}
		}
//		initialization (fixed)
		{
			for (unsigned i = 0; i < N; ++i){
				result[0][i] = 0;
			}
			result[0][0] = 1;
		}

		//main loop
		for (unsigned t = 1; t < shots.size(); ++t){
			T normalization = 0; // scaling factors
			for (unsigned i = 0; i < N; ++i){
				result[t][i] = 0;
				for (unsigned j = 0; j < N; ++j){
					result[t][i] += result[t - 1][j]
							* pow(transition_probabilities(j, i),
									best_path_bias);
				}
				result[t][i] *= pow(
						shooting_distributions[i].probability_density(shots[t]),
						best_path_bias);
				normalization += result[t][i];
			}

			normalization = 1 / normalization;
			for (unsigned i = 0; i < N; ++i){
				result[t][i] *= normalization;
			}
		}

		//Backward
		vector<array<T, N>> backward(shots.size());

		//initialization
		for (unsigned i = 0; i < N; ++i){
			backward[shots.size() - 1][i] = 1.0 / N;
		}
		//main loop
		for (signed t = shots.size() - 2; 0 <= t; --t){
			array<T, N> temp;
			for (unsigned j = 0; j < N; ++j){
				temp[j] = backward[t + 1][j]
						* pow(
								shooting_distributions[j].probability_density(
										shots[t + 1]), best_path_bias);
			}
			T normalization = 0; // scaling factors
			for (unsigned i = 0; i < N; ++i){
				backward[t][i] = 0;
				for (unsigned j = 0; j < N; ++j){
					backward[t][i] += temp[j]
							* pow(transition_probabilities(i, j),
									best_path_bias);
				}
				normalization += backward[t][i];
			}
			normalization = 1 / normalization;
			for (unsigned i = 0; i < N; ++i){
				backward[t][i] *= normalization;
			}
		}

		//fb_combined
		for (unsigned t = 0; t < shots.size(); ++t){
			T normalization = 0;
			for (unsigned i = 0; i < N; ++i){
				result[t][i] *= backward[t][i];
				result[t][i] = pow(result[t][i], best_path_bias);
				normalization += result[t][i];
			}
			normalization = 1 / normalization;
			for (unsigned i = 0; i < N; ++i){
				result[t][i] *= normalization;
			}
		}

//		initialization (fixed)
		{
			for (unsigned i = 0; i < N; ++i){
				result[0][i] = 0;
			}
			result[0][0] = 1;
		}

		return result;
	}

	static vector<Shooter> best_shooters(
			const vector<array<T, N>>& forward_backward){
		vector<Shooter> result(forward_backward.size());
		for (unsigned t = 0; t < forward_backward.size(); ++t){
			Shooter best_shooter = -1;
			T best_forward_backward = -numeric_limits<T>::infinity();
			for (unsigned i = 0; i < N; ++i){
				if (forward_backward[t][i] > best_forward_backward){
					best_shooter = i;
					best_forward_backward = forward_backward[t][i];
				}
			}
			result[t] = best_shooter;
			assert(0 <= best_shooter && best_shooter < N);
		}
		return result;
	}

	vector<Shooter> best_path(const vector<Shot>& shots) const{
		vector<array<Shooter, N>> saved(shots.size());
		array<T, N> probability_temp1;
		array<T, N> probability_temp2;

		for (unsigned i = 0; i < N; ++i){
			probability_temp1[i] = log(
					shooting_distributions[i].probability_density(shots[0]));
			saved[0][i] = -1;
		} // initial state

		//fixed initial
		for (unsigned i = 0; i < N; ++i){
			probability_temp1[i] = -numeric_limits<T>::infinity();
		}
		probability_temp1[0] = 0;
		probability_temp2 = probability_temp1;
		saved[0][0] = 0;

		for (unsigned t = 1; t < shots.size(); ++t){
			for (unsigned i = 0; i < N; ++i){
				Shooter max_state = -1;
				T max_probability = -numeric_limits<T>::infinity();
				for (unsigned j = 0; j < N; ++j){
					if (probability_temp1[j]
							+ log(transition_probabilities(j, i))
							> max_probability){ //Todo what if >= ?
						max_probability = probability_temp1[j]
								+ log(transition_probabilities(j, i));
						max_state = j;
					}
				}
				probability_temp2[i] = max_probability
						+ log(
								shooting_distributions[i].probability_density(
										shots[t]));
				saved[t][i] = max_state;
				assert(0 <= max_state && max_state < N);
			}
			probability_temp1 = probability_temp2;
		}

		//trace back
		vector<Shooter> result(shots.size());
		Shooter max_state = 0;
		T max_probability = -numeric_limits<T>::infinity();

		for (unsigned i = 0; i < N; ++i){
			if (probability_temp2[i] > max_probability){
				max_probability = probability_temp2[i];
				max_state = saved[shots.size() - 1][i];
			}
		}
		result[shots.size() - 1] = max_state;

		for (int t = shots.size() - 2; t >= 0; --t){
			result[t] = saved[t + 1][result[t + 1]];
		}
		return result;
	}

	static unsigned path_similarity(
			const vector<Shooter>& path1,
			const vector<Shooter>& path2){
		assert(path1.size() == path2.size());
		unsigned score = 0;
		for (unsigned i = 0; i < path1.size(); ++i){
			score += (path1[i] == path2[i]);
		}
		return score;
	}

	static vector<T> score_framewise(
			const vector<array<T, N>>& forward_backward
			, const vector<reference_wrapper<vector<Shooter>>>& paths){ //Todo, better to get iterators instead of vectors and const?
				vector<T> result(paths.size(), 0);
				for (unsigned t = 0; t < forward_backward.size(); ++t){ // Todo frame_size-1??
					for (unsigned i=0; i<paths.size(); ++i){
						result[i] += forward_backward[t][paths[i].get()[t]];
					}
				}
				return result;
			}

			static vector<T> score_sequencewise_log(
					const vector<array<T, N>>& forward_backward
					, const vector<reference_wrapper<vector<Shooter>>>& paths){ //Todo, better to get iterators instead of vectors and const?
				vector<T> result(paths.size(), 0);
				for (unsigned t = 0; t < forward_backward.size(); ++t){ // Todo frame_size-1??
					for (unsigned i=0; i<paths.size(); ++i){
						result[i] += log(forward_backward[t][paths[i].get()[t]]);
					}
				}
				return result;
			}

			void save(string file_name){
				save_eigen(transition_probabilities, file_name + "_transition_probabilities");
				for(size_t i = 0; i < shooting_distributions.size(); ++i){
					std::stringstream ss;
					ss << i;
					shooting_distributions[i].save(file_name+"_shooting_distribution" + ss.str());
				}
			}

		};

#endif /* MODEL_H_ */
